package com.bo.day20241119;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * 并查集
 *
 * @Author: gpb
 * @Date: 2024/11/19 16:15
 * @Description:
 */
public class UnionFind {
    public static class Node<V> {
        public V value;

        public Node (V value) {
            this.value = value;
        }
    }

    /**
     * 并查集具体实现
     *
     * @param <V>
     */
    public static class UF<V> {
        /**
         * 节点
         */
        private HashMap<V, Node<V>> nodes;
        /**
         * 父节点
         */
        private HashMap<Node<V>, Node<V>> parents;
        /**
         * 节点的集合
         */
        private HashMap<Node<V>, Integer> sizeMap;

        public UF (List<V> values) {
            nodes = new HashMap<>();
            parents = new HashMap<>();
            sizeMap = new HashMap<>();
            for (V value : values) {
                Node<V> node = new Node<>(value);
                nodes.put(value, node);
                parents.put(node, node);
                sizeMap.put(node, 1);
            }
        }

        /**
         * 判断两个节点是否是同一个集合
         *
         * @param a
         * @param b
         * @return
         */
        public boolean isSameSet (Node a, Node b) {
            return findFather(a) == findFather(b);
        }

        /**
         * 合并
         *
         * @param a
         * @param b
         */
        public void union (Node a, Node b) {
            // 拿到两个节点根节点
            Node aFather = findFather(a);
            Node bFather = findFather(b);
            // 如果两个根节点不一样,需要合并
            if (aFather != bFather) {
                // 小挂大
                int aSize = sizeMap.get(aFather);
                int bSize = sizeMap.get(bFather);
                if (aSize >= bSize) {
                    parents.put(bFather, aFather);
                    sizeMap.put(aFather, aSize + bSize);
                    sizeMap.remove(bFather);
                } else {
                    parents.put(aFather, bFather);
                    sizeMap.put(bFather, aSize + bSize);
                    sizeMap.remove(aFather);
                }
            }
        }

        /**
         * 获取集合的数量
         *
         * @return
         */
        public int size () {
            return sizeMap.size();
        }

        /**
         * 给你一个节点，请你往上到不能再往上，把代表返回
         *
         * @param cur
         * @return
         */
        private Node<V> findFather (Node<V> cur) {
            // 优化点,防止扁平化
            Stack<Node<V>> path = new Stack<>();
            while (cur != parents.get(cur)) {
                cur = parents.get(cur);
                path.push(cur);
            }
            // 扁平化
            while (!path.isEmpty()) {
                parents.put(path.pop(), cur);
            }
            return cur;
        }
    }
}
